# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

import logging
import os
import ujson as json
import json as base_json
import zlib

from coin.flow.subscriber import _FlowSubscriber_SimpleQueue


# in archive: record.data is bytes form, need to be parsed first,
# in live: record.data is already dict or list
def parse_topic_record_data_from_live(record):
  assert isinstance(record.data, (dict, list))
  return record.data


def parse_topic_record_data_from_archive(record):
  assert hasattr(record.data, "tobytes")
  message = json.loads(record.data.tobytes())
  return message


# when integer exceeds int64: only base_json can load
def parse_topic_record_data_from_archive_base_json(record):
  assert hasattr(record.data, "tobytes")
  message = base_json.loads(record.data.tobytes())
  return message


def gdax_onbook_hash(message):
  changes_or_message = message.get('changes', message)
  return (message['product_id'], zlib.crc32(json.dumps(changes_or_message).encode()))


def gdax_onfull_hash(message):
  if message['type'] != 'match':
    return zlib.crc32(json.dumps(message).encode())
  return (message['type'], message['product_id'], message['sequence'])


def huobi_spot_depth_hash(message):
  return (message['ch'], message['tick']['ts'])


def huobi_spot_delta_depth_hash(message):
  return (message['ch'], message['ts'])


def huobi_spot_trade_hash(message):
  return (message['ch'], message['tick']['id'])


def huobi_depth_hash(message):
  return (message['ch'], message['tick']['mrid'])


def huobi_trade_hash(message):
  return (message['ch'], message['tick']['id'])


def okex_diff_feed_hash(message):
  # Same diff update in different stream may have different exchange timestamp.
  try:
    del message['data']['timestamp']
  except KeyError:
    # Ignore key error.
    pass
  data = json.dumps(message['data'])
  return zlib.crc32(data.encode())


def data_hash(message):
  data = json.dumps(message['data'])
  return zlib.crc32(data.encode())


okex_feed_hash = data_hash


def message_hash(message):
  data = json.dumps(message)
  return zlib.crc32(data.encode())


def bitmex_trdmatchid(message):
  return tuple([zlib.crc32(each_trd['trdMatchID'].encode()) for each_trd in message['data']])


def bitflyer_feed_hash(message):
  if 'C' not in message:
    if len(message) > 0:
      # this doesn't happen.
      logging.warning("non-null message doesn't have 'C': %s", message)
      return data_hash(message)
    else:
      return None
  return message['C']


def binance_ws_book_hash(message):
  if 'data' in message:
    message = message['data']
  return (message['s'], message['u'], message['U'])


def binance_ws_trade_hash(message):
  if 'data' in message:
    message = message['data']
  return (message['s'], message['t'])


def binance_ws_book20_hash(message):
  return (message['stream'], message['data']['lastUpdateId'])


def get_hash(hash_func, flow_sub, robust_json_parser=False):
  if robust_json_parser:
    archive_parser = parse_topic_record_data_from_archive_base_json
  else:
    archive_parser = parse_topic_record_data_from_archive

  if isinstance(flow_sub, _FlowSubscriber_SimpleQueue):
    return lambda record: hash_func(parse_topic_record_data_from_live(record))
  else:
    if ('RUST_FEED_HANDLER' in os.environ) and os.environ['RUST_FEED_HANDLER'] == '1':
      from native_rust import feed_handler
      if hash_func == data_hash:
        print("Use rust topic_record_data_hash")
        return feed_handler.topic_record_data_hash
      elif hash_func == message_hash:
        print("Use rust topic_record_data_hash")
        return feed_handler.topic_record_data_hash
      elif hash_func == huobi_depth_hash:
        print("Use rust huobi_depth_hash")
        return feed_handler.huobi_depth_hash
      elif hash_func == huobi_trade_hash:
        print("Use rust huobi_trade_hash")
        return feed_handler.huobi_trade_hash
      elif hash_func == huobi_spot_depth_hash:
        return feed_handler.huobi_spot_hash
      else:
        print("Failed to use rust hash.")
    return lambda record: hash_func(archive_parser(record))


# from native_rust import topic_record
# from native_rust import feed_handler
# import simdjson

# def huobi_depth_hash(record):
#   if ('RUST_HASH' in os.environ) and os.environ['RUST_HASH'] == '1':
#     b = record.data.tobytes()
#     # idx = b.find(b'mrid')
#     idx = 67
#     return b[15:21], b[idx:idx + 18]

#     # pj = simdjson.ParsedJson(record.data.tobytes())
#     # return (pj.items('.ch'), pj.items('.tick.mrid'))
#   else:
#     message = parse_topic_record_data(record)
#     return (message['ch'], message['tick']['mrid'])

# def huobi_trade_hash(record):
#   if ('RUST_HASH' in os.environ) and os.environ['RUST_HASH'] == '1':
#     pj = simdjson.ParsedJson(record.data.tobytes())
#     return (pj.items('.ch'), pj.items('.tick.id'))
#   else:
#     message = parse_topic_record_data(record)
#     return (message['ch'], message['tick']['id'])
